Skip to content

feat(recipes): Kimi-K3 PD on SGLang, selectable alongside the vLLM one - #71

Merged
jiejingzhangamd merged 8 commits into
mainfrom
feat/kimi-k3-pd
Aug 2, 2026
Merged

feat(recipes): Kimi-K3 PD on SGLang, selectable alongside the vLLM one#71
jiejingzhangamd merged 8 commits into
mainfrom
feat/kimi-k3-pd

Conversation

@jiejingzhangamd

Copy link
Copy Markdown
Collaborator

Kimi-K3 PD had one shape (vLLM, MooncakeConnector). This adds the SGLang shape and makes the choice explicit, because the two disaggregate by different mechanisms rather than being two spellings of the same thing:

Engine Mechanism Manifest Image
vLLM --kv-transfer-config with a MooncakeConnector pd/deploy.yaml vllm/vllm-openai-rocm:kimi-k3
SGLang --disaggregation-mode + Mooncake bootstrap handshake pd-sglang/deploy.yaml lmsysorg/sglang-rocm:rocm720-mi35x-k3-20260727

The manual page now presents them as a choice inside the PD tab.

The image is the load-bearing detail

The general lmsysorg/sglang tags carry no Kimi-K3 support at all — no kimi_k3.py, nothing in the registry. Only that dated -k3- build has it:

/sgl-workspace/sglang/python/sglang/srt/models/
    kimi_k3.py   kimi_k3_vl.py   kimi_k25.py

class KimiK3ForConditionalGeneration
class KimiK3LinearForCausalLM

Verifying that is less obvious than it looks, and I got it wrong twice before getting it right. The image installs sglang as an editable package, so import sglang does not resolve to the source tree, and a registry probe returns an empty list — doubly so because rocminfo fails in a GPU-less container and some model imports abort. Both of my first two checks said "no Kimi support", which would have been the wrong conclusion to act on. The files are under /sgl-workspace/sglang/python/.

Weights come from NFS, and that is not a preference

Kimi-K3 is ~1.5 TB; the decode node here has 147 GB free. Both nodes seeing one path is the only option, so the manifest points at NFS rather than a node-local directory.

Validation status — stated, not implied

Neither PD combination has been run for this model. The SGLang PD path is validated cross-node (the Qwen run in the PD-fixes PR: prefill chi2800 → decode chi2866, KV over RoCE, correct answer), but that was Qwen3-0.6B at TP1 — not Kimi-K3's size or settings. The vLLM PD path is unproven on this fabric entirely.

The status table says exactly that rather than letting the model inherit a validation it did not earn.

Depends on

The SGLang PD path needs the five fixes in the cross-node PD PR — in particular the advertise-host fix, without which the router cannot pair prefill with decode, and the Mooncake HIP gate, without which the KV transfer dies in hipIpcOpenMemHandle.

🤖 Generated with Claude Code

https://claude.ai/code/session_01Pw7JSvdQb796xh5pEcctLz

Kimi-K3 PD had one shape (vLLM, MooncakeConnector). Adds the SGLang shape and
makes the choice explicit in the manual, because the two disaggregate by different
mechanisms rather than being two spellings of one thing:

    vLLM     --kv-transfer-config with a MooncakeConnector
    SGLang   --disaggregation-mode + a Mooncake bootstrap handshake

The image matters more than usual here. The general `lmsysorg/sglang` tags carry
no Kimi-K3 support at all — no kimi_k3.py, nothing in the registry — while
`lmsysorg/sglang-rocm:rocm720-mi35x-k3-20260727` has KimiK3ForConditionalGeneration,
KimiK3LinearForCausalLM and kimi_k3_vl. Checking that is less obvious than it
sounds: the image installs sglang as an editable package, so `import sglang` does
not resolve to the source tree and a registry probe reports an empty list, doubly
so because rocminfo fails in a GPU-less container. The files are under
/sgl-workspace/sglang/python/sglang/srt/models/.

Weights come from NFS rather than a node-local path: Kimi-K3 is ~1.5 TB and the
decode node here has 147 GB free, so both nodes seeing one path is the only
option, not a preference.

Neither PD combination has been run for this model yet — the SGLang PD path itself
is validated cross-node (see the Qwen run in the PD fixes), but not with Kimi-K3's
size or settings, and the vLLM PD path is unproven on this fabric entirely. The
status table says so rather than implying the model inherits that validation.

Signed-off-by: Zhang, Jiejing <jiejing.zhang@amd.com>
infera's fp8 KV-cache default assigned `server_args.kv_cache_dtype` after
`ServerArgs.from_cli_args()`. Newer SGLang freezes server_args once resolved:

    AttributeError: server_args.kv_cache_dtype assigned after resolution;
    server_args is read-only -- use get_context().override(source, ...)

which crash-looped every worker on lmsysorg/sglang-rocm:rocm720-mi35x-k3-20260727,
the only SGLang build carrying Kimi-K3.

Set it on the parsed args before resolution instead. That needs no override API
and works on both the old and the new SGLang, where calling get_context().override
would only work on the new one.

Also carries the two mounts the PD recipe needs and had lost by being derived from
main rather than from the PD-fixes branch: hostNetwork on the workers, and the host
libionic. Without them the recipe reproduces the exact failure that branch fixes —
0 RDMA devices, then Mooncake down its HIP-IPC path. Verified in this run that the
injection fires: "infera-exec: replaced .../libionic.so.1.0.54.0-149 with the host
libionic build".

Signed-off-by: Zhang, Jiejing <jiejing.zhang@amd.com>
Both workers hardcoded a 1800 s wait for the engine's /health. That is generous
for local NVMe and impossible for NFS: Kimi-K3 read its 96 shards in 502 s from
local disk and ~95 min from NFS with both PD nodes competing for the same mount.

The failure is worse than a slow start. The worker times out mid-load, exits 1,
restarts, and begins the load again from zero — so it never finishes, and the pod
looks like a crash loop rather than "the storage is too slow for this deadline":

    TimeoutError: SGLang not ready after 1800s

Load time tracks the storage, not the model, so this cannot have one right value
baked in. INFERA_ENGINE_READY_TIMEOUT now sets it, on both the SGLang and vLLM
workers, defaulting to the previous 1800 and falling back to it on a malformed
value rather than raising during startup.

The Kimi-K3 PD recipe sets 7200, with the measured numbers next to it so the next
person can tell whether their storage needs more.

Signed-off-by: Zhang, Jiejing <jiejing.zhang@amd.com>
Both PD manifests were derived from ancestors that predate what running PD
actually taught us, and each inherited the same gaps. Bringing them level:

  hostNetwork on the workers      the RDMA rails are host interfaces; the flannel
                                  pod network cannot reach them
  host libionic mount             the vendor base's provider must match the host
                                  ionic kernel ABI or libibverbs sees no devices
  /dev/infiniband + privileged    RDMA needs the device nodes and pinned memory
  INFERA_ENGINE_READY_TIMEOUT     1800 s is impossible on NFS-backed weights

The SGLang manifest additionally carries what its own bring-up cost: memory 96Gi
(a Qwen-0.6B number) against 1.5 TB of shards, a mamba state cache sized from a
negative residual, and the mem-fraction trap documented inline.

None of this is speculative — every line replaces a failure whose message pointed
somewhere other than its cause: "0 RDMA devices" for an ABI mismatch, a crash loop
for a timeout, a bare "-9" for two different problems in turn.

Signed-off-by: Zhang, Jiejing <jiejing.zhang@amd.com>
The overlay could ship code the vendor base lacks, but not fix code it has. That
gap became blocking: Kimi-K3 PD on vLLM dies during KV registration because the
Mooncake connector destructures every Mamba-family layer as a 2-tuple —

    if isinstance(layer_spec, MambaSpec):
        conv, _ = cache_or_caches

which is Mamba2's shape. Kimi-K3 uses KDA linear attention with a different
number of state tensors, so every rank raises "too many values to unpack".
MambaSpec.shapes is already a variable-length tuple, so the arity assumption is
the connector's, not the spec's.

The obvious fix — bake the patch into Dockerfile.vllm — reintroduces exactly the
fork this overlay exists to avoid, and for a 42 GB image. So the payload now
carries deploy/docker/patches/vllm/ and infera-exec applies them before exec.
Each script locates its target through the installed package, is idempotent, and
no-ops when its anchor is gone, so a base that already carries a fix is untouched
and the patches survive a base bump.

Guarded rather than unconditional: skipped when vllm is not importable, skipped
by INFERA_SKIP_ENGINE_PATCHES, and a failing patch warns instead of aborting —
a mis-anchored patch must not take down a worker that would otherwise serve.

Verified in a stock vllm/vllm-openai-rocm:kimi-k3: the mamba unpack is applied,
the old line is gone, and the pre-existing blocksize/prom patches apply too — so
this also closes the gap where an overlay deployment silently lacked fixes the
baked image had.

Signed-off-by: Zhang, Jiejing <jiejing.zhang@amd.com>
…s known

infera resolves the Pod IP under Kubernetes discovery and logs it —

    k8s discovery: advertising Pod IP 137.220.61.177

— but disagg_meta, which carries Mooncake's bootstrap_addr, was already built
during arg parsing, while advertise_host was still unset. So the address froze as
the bind host and the prefiller published:

    bootstrap_addr: 'http://0.0.0.0:8998'

The decoder reads that from kv_transfer_params (vLLM's own comment says decoders
must) and cannot connect:

    Failed to connect to bootstrap server http://0.0.0.0:8998
    Failed to find remote engine_id ...

Both workers stay healthy and registered; the request simply hangs until the
connector's 480 s timeout. Nothing in that path names the address.

This is the same defect as the SGLang KV-endpoint one, in a second place: an
address is derived before the host that makes it routable is known. Rebuild it
after the resolution, touching only the 0.0.0.0 form so an explicitly configured
address is left alone.

The manifest also gains --data-parallel-address $(POD_IP): that fixes the OTHER
half, where the prefiller computes its own address via
get_mooncake_bootstrap_addr() -> parallel_config.data_parallel_master_ip. Both
paths have to be right; fixing either alone leaves PD hanging.

With both, Kimi-K3 PD answers across two nodes: prefill on chi2800, decode on
chi2866, 1.5 TB MXFP4 at TP8 from local NVMe, "The capital of France is Paris."

Signed-off-by: Zhang, Jiejing <jiejing.zhang@amd.com>
v0.2.1 predates everything cross-node PD needed. Comparing the published image
against the current build:

                        v0.2.1   v0.2.2
    libionic injection      no      yes
    runtime engine patches  no      yes   (8 patch scripts)
    native trees            2       2

So anyone following a recipe today gets an overlay that cannot bring up PD: no
libionic swap means libibverbs rejects every ionic device, and no runtime patches
means the Mooncake connector still destructures Kimi-K3's Mamba layers as a
2-tuple. Both were fixed in this branch and neither was reachable.

inferaimage/infera-overlay:v0.2.2 is public and pulls with no credentials —
verified from a clean DOCKER_CONFIG, since "we pushed it" and "a reader can get
it" have already turned out to be different things twice in this repo.

Validated with this image: Kimi-K3 PD across chi2800 and chi2866, 1.5 TB MXFP4 at
TP8 from local NVMe, ~6 min to both workers registered, 8 RDMA devices visible to
libibverbs on each side, and a correct answer through the router.

Signed-off-by: Zhang, Jiejing <jiejing.zhang@amd.com>
Two problems, one fix.

The recipes still carried 20 references to rocm/infera-overlay:latest — the name
I built under locally while writing #64, which 404s on Docker Hub. #67 fixed those
on a branch cut from main; this branch predates that merge and had them back. A
name that only exists on one laptop keeps coming back because a tag is a promise
nobody checks.

And the tags that do resolve are mutable. vllm/vllm-openai-rocm:kimi-k3 was
rebuilt on 2026-07-27 during this very session, which is what invalidated the
"prefix caching fails engine init" note: the tag was the same and the image was
not. A recipe that says "this was validated" has to name the thing that was
validated.

So every image is now pinned by digest:

    inferaimage/infera-overlay          @6918eff3  (v0.2.2, this session's fixes)
    vllm/vllm-openai-rocm:kimi-k3       @5aa7e626
    lmsysorg/sglang-rocm:...k3-20260727 @3c01f73f
    lmsysorg/sglang:v0.5.15...mi35x     @40e940a0

Zero mutable tags remain under examples/recipes/. Each digest is one I pulled and
ran, not one read off a page. All nine manifests still pass
`kubectl apply --dry-run=server` against the live CRD.

The readable tag is kept alongside the digest so the ref still says what it is.

Signed-off-by: Zhang, Jiejing <jiejing.zhang@amd.com>
@jiejingzhangamd
jiejingzhangamd merged commit 2aeb17c into main Aug 2, 2026
28 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant